# -*_ coding=utf-8 _*_
# cnn_eye_cls
# 2023/6/28
# Project:1.py
# Creater:Administrator
# Create time:2023-06-28-16-26
# IDE:PyCharm

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision.models as models
import torchvision.datasets as datasets
import torchvision.transforms as transforms
import numpy as np
import matplotlib.pyplot as plt
import os

os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'

# 数据预处理
transform = transforms.Compose([  # 一系列的图像变换
    transforms.Resize((224, 224)),  # 将图像转化为224*224
    transforms.ToTensor(),  # 将图像转化为tensor
    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])  # 归一化
])

transform_test = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])

BATCH_SIZE = 64
device = 'cuda' if torch.cuda.is_available() else 'cpu'

dataset_train = datasets.ImageFolder(r'D:\\coding\\jupyterFILE\\Graph\\gnnlstm\\eye_data\\train', transform)
dataset_test = datasets.ImageFolder(r'D:\\coding\\jupyterFILE\\Graph\\gnnlstm\\eye_data\\test', transform_test)
# 读取数据
print(len(dataset_train.imgs))
print(len(dataset_test.imgs))

# 导入数据
train_loader = torch.utils.data.DataLoader(dataset_train, batch_size=BATCH_SIZE, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset_test, batch_size=BATCH_SIZE, shuffle=False)

model_ft = models.resnet18(pretrained=True)

for param in model_ft.parameters():
    # 屏蔽预训练模型的权重,只训练最后一层的全连接的权重。
    param.requires_grad = False

# 修改最后一层的全连接层
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, 4)

model_ft = model_ft.to(device)

optimizer = optim.Adam(model_ft.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()

acc_train = []
acc_test = []
x_val = []


# 训练模型
def train(model, device, train_loader, optimizer, epoch):
    model.train()
    correct = 0
    for batch_idx, (data, target) in enumerate(train_loader):  # batch_idx批次索引，data数据，target标签
        data, target = data.to(device), target.to(device)  # 将数据放在device上
        optimizer.zero_grad()  # 梯度置零
        output = model(data)  # 前向传播
        loss = criterion(output, target)  # 计算损失
        loss.backward()  # 反向传播
        optimizer.step()  # 更新参数
        pred = output.max(1, keepdim=True)[1]  # 找到概率最大的下标
        correct += pred.eq(target.view_as(pred)).sum().item()  # 计算正确的数量

        if batch_idx % 20 == 0:
            print('Train Epoch:{} [{}/{} ({:.0f}%)]\tLoss:{:.6f}'.format(
                epoch + 1, batch_idx * len(data), len(train_loader.dataset), 100. * batch_idx / len(train_loader),
                loss.item()
            ))
    acc_train.append(100. * correct / len(train_loader.dataset))
    print('\nTrain Accuracy:{}/{} ({:.2f}%)'.format(
        correct, len(train_loader.dataset), acc_train[-1])
    )  # 计算训练集的准确率


# 测试模型
def test(model, device, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():  # 不需要计算梯度
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += criterion(output, target).item()  # sum up batch loss
            pred = output.max(1, keepdim=True)[1]  # 找到概率最大的下标
            correct += pred.eq(target.view_as(pred)).sum().item()  # 计算正确的数量

    test_loss /= len(test_loader.dataset)
    acc_test.append(100. * correct / len(test_loader.dataset))
    print('\nTest set: Average loss:{:.4f}, Accuracy:{}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset), 100. * correct / len(test_loader.dataset)
    ))


plt.ion()
for epoch in range(100):
    train(model_ft, device, train_loader, optimizer, epoch)
    test(model_ft, device, test_loader)

    x_val.append(epoch)

    plt.xlim(0, 100)
    plt.ylim(0, 1)
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy')
    plt.title('Training and validation accuracy')

    # plt.scatter(epoch, acc_train[-1], c='red')
    # plt.scatter(epoch, acc_test[-1], c='blue')

    plt.plot(x_val, acc_train, 'r')
    plt.plot(x_val, acc_test, 'b')

    plt.legend(['Training Acc', 'Validation Acc'], loc='lower right')
    plt.draw()
    plt.pause(0.01)

plt.ioff()
